from numbers import Number
import tensorflow as tf
import os
import numpy as np
import matplotlib.pyplot as plt


def himmelblau(x: Number, y: Number):
    """
    计算 himmelblau 函数的值
    :param x:
    :param y:
    :return: 计算后的函数值
    """
    return (x ** 2 + y - 11) ** 2 + (x + y ** 2 - 7) ** 2


if __name__ == '__main__':
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
    x = np.arange(-6, 6, 0.1)
    y = np.arange(-6, 6, 0.1)

    # print(x)
    # print(y)
    print(x.shape)
    print(y.shape)

    X, Y = np.meshgrid(x, y)
    print(X.shape)
    # print(Y.shape)
    # plt.plot(X, Y, marker='o', color='k', linestyle='none')
    # plt.show()

    # print(X)
    Z = himmelblau(X, Y)
    # print(Z.shape)
    # print(Z)
    figure = plt.figure("himmelblau")
    ax = figure.add_subplot(projection='3d')
    ax.plot_surface(X=X, Y=Y, Z=Z)
    ax.view_init(60, -30)
    ax.set_xlabel('X')
    ax.set_ylabel('Y')
    ax.set_zlabel('Z')
    # plt.show()

    x = tf.constant([4.0])
    y = tf.constant([0.0])

    for step in range(200):
        with tf.GradientTape() as tape:
            tape.watch([x, y])
            z = himmelblau(x, y)

        grads = tape.gradient(z, [x, y])
        print(grads)